#!/usr/bin/perl
use strict;
use warnings;
use Bio::SeqIO;
use Bio::Seq;
use Getopt::Long;
use Data::Dumper;
#standing question: if output is 21-mer, which 22-mer does it correspond to?

#print "version x\n";
#user-determined parameters
my $shrna_length = 22;
my $svm_oligo_length = 50; 	# length of shRNA + flanks to feed into previously trained svm in R
my $guide_3p_flank = 11;  #how many bases in the svm_oligo 3' on the guide strand (5' on the mRNA)?
my $guide_5p_flank = $svm_oligo_length - $shrna_length - $guide_3p_flank; # 5' to guide, 3' on mRNA
my $sfold_len = 0;


#set up usage message for errors
my $usage = qq{
Usage: screen_shrna.pl --gene=<genename> --fasta=<in.fasta> --proden=<in.proden> --contra=<in.posteriors> --mirscan=<in.mir> --sirna_sfold=<in.sirna> --out=<outfile>

Note: 	all files are required
		output from contrafold should be in posteriors format
		output from sfold sirna should be sirna.out
		
};

#load in filenames and lists
my $gene_name = "";
my $fasta_file = "";  #gene sequence
my $proden_file = "";
my $contra_file = "";  #contrafold output in posterior format
my $mirscan_file = "";
my $broad_file = "";  #two columns, start site and guide strand sequence
my $sirna_sfold_file = "";
my $out_file = "output.txt";
my $result = GetOptions (	"gene=s" => \$gene_name,				# name of gene, if no proden file
							"fasta=s" => \$fasta_file, 				# infile
							"proden=s" => \$proden_file,			# if present, the proden field will be populated for training
 							"contra=s" => \$contra_file, 			# secondary structure from contrafold in posterior format
							"mirscan=s" => \$mirscan_file, 			# scores from mirscan
#							"broad=s" => \$broad_file, 				# list of good oligos from the broad
							"sirna_sfold=s" => \$sirna_sfold_file, 	# output from sfold (2 lines per oligo) with leading spaces removed
							"out=s" => \$out_file 					# output
						);
#test for presence of required information
die ($usage) if ($fasta_file eq "");

#read in first (and only) sequence- we do this one gene at a time...
my $gene_obj = Bio::SeqIO->new('-file' => $fasta_file,
                         '-format' => "fasta");
my $gene = $gene_obj->next_seq();
my $gene_seq = $gene->seq();
$gene_seq =~ tr/AGCTUagctu/AGCTTAGCTT/;  #we're going with T and upper case here
my $gene_rc = reverse($gene_seq);
$gene_rc =~ tr/AGCT/TCGA/; #already switched T to U and upper case...
my $gene_length = length($gene_seq);
if ($gene_name eq "") {
	$gene_name = $gene->id();
	}

print "gene id: $gene_name \tlength: $gene_length\nReading other files...\n";

#read in output from contrafold, mirscan, broad, sirna
open(my $contra_fh,"<",$contra_file) || die ("no contrafile\n");
open(my $mirscan_fh,"<",$mirscan_file) || die ("no mirscan file\n");
#open(my $broad_fh,"<",$broad_file);
open(my $sirna_sfold_fh,"<",$sirna_sfold_file) || die ("no sirna file\n");

#open output file
open(my $out_fh,">",$out_file);

#open and parse proden: gene \t location \t guideseq \t proden
my %proden_scores;
if ($proden_file ne "") {
	open(my $proden_fh,"<",$proden_file);
	while (my $proden_line = <$proden_fh>) {
		chomp($proden_line);
		my @prodens = split(/\s+/,$proden_line);
		my $site = $prodens[1];
		my $doublecheck = substr($gene_seq,$site-1,22);
		$doublecheck = revcomp($doublecheck);
		if ($prodens[4] ne $doublecheck) {
			die ("at location $site, $prodens[4] ne $doublecheck\nis your proden from the right gene?\n");
			}
		$proden_scores{$site} = $prodens[2]; #set the proden score for that site to the score
#		print"assigning proden: $site, $proden_scores{$site}\n";
		}
	}

my @header; #put in headers here...
push (@header,"ID");
if ($proden_file ne "") {
		push (@header,"Score");
		} 
for (my $i=0;$i<$svm_oligo_length;$i++) {
	push (@header, "a$i","b$i","c$i");
	}
for (my $i=0;$i<$svm_oligo_length;$i++) {
	push (@header, "pair$i");
	}
push (@header,"mirscan");

#commented out because non-training samples must not have a score column
#else {
#for (my $i = 0; $i < (length($gene_seq)-21);$i++) {
#		$proden_scores{$i} = "NA";
#		}
#	}
#parse contra
my @contra_probs;
for (my $i = 0; $i < (length($gene_seq));$i++) {
	my $contra_line = <$contra_fh>;
	chomp($contra_line);
	my $num_pairings = ($contra_line =~ tr/:/:/);
	my @contras = split(/\s+/,$contra_line);
	$contras[1] =~ tr/Ttuagc/UUUAGC/;
	my $doublecheck = substr($gene_seq,$i,1);
	$doublecheck =~ tr/Ttuagc/UUUAGC/;
	if ($contras[1] ne $doublecheck) {
		die ("$contras[1] ne $doublecheck\nis your contrafold from the right gene?\n");
		}
	if ($num_pairings == 0) {
		unless (exists($contra_probs[$i])) {
			$contra_probs[$i] = 0;
			}
		}
	while ($num_pairings > 0) {
		my $pairing = pop(@contras);
		my @pair_info = split(":",$pairing);
		$contra_probs[$i] += $pair_info[1];
		$contra_probs[($pair_info[0]-1)] += $pair_info[1];
		$num_pairings--;
		}
	}
#parse mirscan
my @mirscan_scores;
for (my $i = 0; $i < (length($gene_seq)-21);$i++) {
	my $mirscan_line = <$mirscan_fh>;
	chomp($mirscan_line);
	my @mirscans = split(/\s+/,$mirscan_line);
	my $site=$mirscans[0];
	my $doublecheck = substr($gene_seq,$i,22);
	$doublecheck =~ tr/AGCTUagctu/AGCUUAGCUU/;  #mirscan outputs U's
	if ($mirscans[2] ne $doublecheck) {
		die ("at $i, $mirscans[2] ne $doublecheck\nis your mirscan from the right gene?\n");
		}
	$mirscan_scores[$site] = $mirscans[4];
	}
#parse broad: columns, 1based_match_pos=7,adj_score=12,target=13
# my %broad_oligo_scores;
# BROAD: while (my $broad_line = <$broad_fh>) {
# 	chomp($broad_line);
# 	my @broads = split(/\s+/,$broad_line);
# 	if ($broads[0] =~ m/Input/) {
# 		next BROAD;
# 		}
# 	my $doublecheck = substr($gene_seq,($broads[7]-1),22);
# 	$doublecheck =~ tr/AGCTUagctu/AGCTTAGCTT/; #broad outputs T's
# 	if ($broads[13] ne $doublecheck) {
# 		die ("is your broad from the right gene?\n");
# 		}
# 	$broad_oligo_scores{($broads[7]-1)} = $broads[12];
# 	}
#parse sirna (from sfold on wadsworth): two lines per oligo
my %sfold_oligo_scores;
until (<$sirna_sfold_fh> =~ m/-------------/) {
	}
while (my $sfold_line1 = <$sirna_sfold_fh>) {
	my $sfold_line2 = <$sirna_sfold_fh>;
	chomp($sfold_line1);
	chomp($sfold_line2);
	$sfold_line1 =~ s/^\s+//;
	$sfold_line2 =~ s/^\s+//;
	my @sfolds1 = split(/\s+/,$sfold_line1);
	my @sfolds2 = split(/\s+/,$sfold_line2);
#	print "sfold value: @sfolds2\n";
	if ($sfold_len == 0) {
		$sfold_len = scalar(@sfolds2);
		}
	my $site = $sfolds1[1]-21;
	my $doublecheck_gene = substr($gene_seq,$site+2,19);
	$doublecheck_gene =~ tr/AGCTUagctu/AGCUUAGCUU/; #sfold outputs whatever
	my $doublecheck_sfold = substr($sfolds1[2],0,19);
	$doublecheck_sfold =~ tr/AGCTUagctu/AGCUUAGCUU/; #sfold outputs whatever
	if ($doublecheck_gene ne $doublecheck_sfold) {
		die ("at $site, $doublecheck_gene ne $doublecheck_sfold\nis your sfold from the right gene?\n");
		}
	$sfold_oligo_scores{$site} = \@sfolds2;   #the whole line!
	}
close $sirna_sfold_fh;
#final header addition
for (my $i=0;$i<$sfold_len;$i++) {
	push (@header, "sfold$i");
	}
#print "@header\n";
my $hlen = scalar(@header);
if ($proden_file eq "") {
print $out_fh "$header[0]\t";
	for (my $j = 1;$j<$hlen;$j++) {
		print $out_fh "\t$header[$j]";
		}
	print $out_fh "\n";
	}
#######################################################################################
#OK, we have our files parsed, lets make sense of it all
my @gene_oligos;
my $padding = "NNNNNNNNNNNNNNNNNNNNNNNNNNNNNNNNNNN";
BUSINESS: for (my $position = 0; $position < ($gene_length-$shrna_length);$position++) {
	if ($proden_file ne "") { #skip anything missing a proden score for training sets
		unless (exists($proden_scores{$position})) {
			next BUSINESS;
			}
		}
	my $guideseq = revcomp(substr($gene_seq,$position-1,22));
	my $id = $gene_name . "_" . $position . "_" . $guideseq;
	#put oligo with flanks into gene_oligos array
	my $what = $gene_length - $shrna_length - $guide_3p_flank;
	if ($position < ($guide_3p_flank)) {
		my $overhang = $guide_3p_flank - $position;
		$gene_oligos[$position] = substr($padding,0,$overhang) . substr($gene_seq,0,($svm_oligo_length-$overhang));
		}
	elsif ($position > ($gene_length - $shrna_length - $guide_5p_flank)) {
		my $overhang = $position + $shrna_length + $guide_5p_flank - length($gene_seq);
		$gene_oligos[$position] = substr($gene_seq,($position-$guide_3p_flank),($svm_oligo_length-$overhang)) . substr($padding,0,$overhang);
		}
	else {
		$gene_oligos[$position] = substr($gene_seq,$position-$guide_3p_flank,$svm_oligo_length);
		}
	my @vector;  #this holds a row for the output file
	#vector fields: oligoid, prodenscore,150xencodedoligo+flanks,50xcf_pairscore_oligo+flanks,mirscan,sfoldarray
	#put in oligo id and guide strand sequence with flanks
#	$gene_oligos[$position] = revcomp($gene_oligos[$position]);
	push(@vector,$id);
	#put in proden score
	if (exists($proden_scores{$position})) {
		push(@vector,$proden_scores{$position});
		} 
	#load Sciabola's encoding for single bases x 150
	for (my $j = 0; $j < $svm_oligo_length; $j++) {
		my $base = substr($gene_oligos[$position],$j,1);
		if ($base eq "A") {
			push(@vector,-1);
			push(@vector,-1);
			push(@vector,+1);
			}
		elsif ($base eq "C") {
			push(@vector,+1);
			push(@vector,-1);
			push(@vector,-1);
			}
		elsif ($base eq "G") {
			push(@vector,-1);
			push(@vector,+1);
			push(@vector,-1);
			}
		elsif ($base eq "T") {
			push(@vector,+1);
			push(@vector,+1);
			push(@vector,+1);
			}
		else {
			push(@vector,"NA");
			push(@vector,"NA");
			push(@vector,"NA");
			}
		}
	#add in parameters from other sources
	#contrafold pairing probabilities x 50
	for (my $j = $svm_oligo_length; $j > 0 ; $j--) {
		my $place = $position+$j-$guide_3p_flank-1;
		my $vec = scalar(@vector);
#		print "at vec $vec, for site: $place, prob= $contra_probs[$place]\n";
		if (exists($contra_probs[$place])) {
			push(@vector,$contra_probs[$place]);
			}
		else {
			if ($place < $gene_length) {
				print "contra really shouldn't be NA at site $place, gene $gene_name when $place < $gene_length \n";
				}
			push(@vector,"NA");
			}
		}
	#mirscan
	my $len = scalar(@vector);	
#	print "mirscan: setting $len equal to $mirscan_scores[$position]";
	if (exists($mirscan_scores[$position])) {
		push(@vector,$mirscan_scores[$position]);
		}
	else {
		push(@vector,"NA");
		}
#	if (exists($broad_oligo_scores{$position})) {
#		$vector[$other+2] = $broad_oligo_scores{$position};
#		}
#	else {
#		$vector[$other+2] = 0;
#		}
	
	if (exists($sfold_oligo_scores{$position})) {
		for (my $j=0;$j<$sfold_len;$j++) {
			my $sfold_param = $sfold_oligo_scores{$position}->[$j];
			if ($sfold_param =~ m/%/) {
				$sfold_param =~ s/^([0-9][0-9])\.([0-9]{1,2})%/0.$1$2/;
				$sfold_param =~ s/^([0-9])\.([0-9]{1,2})%/0.0$1$2/;
				$sfold_param =~ s/^\*{3,4}%/1.00/;
				}
#			print "sfold parms:$j:$sfold_param\t";
			push(@vector,$sfold_param);
			}
		}
	else {
		for (my $j=0;$j<$sfold_len;$j++) {
			push(@vector,"NA");
			}
		}
	#output to file
	$len = scalar(@vector);
#	print "for oligo $position, vector length: $len, header length: $hlen, slen: $sfold_len\n";
	print $out_fh "$vector[0]";
	for (my $j = 1;$j<$len;$j++) {
		print $out_fh "\t$vector[$j]";
#		print "$vector[$j]\t";
		}
	print $out_fh "\n";
#	print "\n";
	}

sub revcomp {
		my $sequence = $_[0];
		$sequence =~ tr/AGCTUNagctun/TCGAANTCGAAN/;  #reverse complement, use T's
		$sequence = reverse($sequence); 
		return $sequence;
		}
